package com.wistech.solr.elevation;

import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Map;
import java.util.WeakHashMap;

import javax.xml.xpath.XPath;
import javax.xml.xpath.XPathConstants;
import javax.xml.xpath.XPathExpressionException;
import javax.xml.xpath.XPathFactory;

import org.apache.lucene.index.IndexReader;
import org.apache.solr.cloud.ZkController;
import org.apache.solr.common.SolrException;
import org.apache.solr.core.Config;
import org.apache.solr.core.SolrCore;
import org.apache.solr.search.SolrIndexSearcher;
import org.apache.solr.util.DOMUtil;
import org.apache.solr.util.RefCounted;
import org.apache.solr.util.VersionedFile;
import org.apache.solr.util.plugin.SolrCoreAware;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.w3c.dom.Node;
import org.w3c.dom.NodeList;
import org.xml.sax.InputSource;

/**
 * A component to elevate some documents to the top of the result set.
 * 
 * 
 */
public class FileQueryElevationComponent extends
		AbstractQueryElevationComponent implements SolrCoreAware {
	private static Logger log = LoggerFactory
			.getLogger(FileQueryElevationComponent.class);
	static final String CONFIG_FILE = "config-file";
	// For each IndexReader, keep a query->elevation map
	// When the configuration is loaded from the data directory.
	// The key is null if loaded from the config directory, and
	// is never re-loaded.
	private final Map<IndexReader, Map<String, ElevationBean>> elevationCache = new WeakHashMap<IndexReader, Map<String, ElevationBean>>();

	@Override
	public void inform(SolrCore core) {
		super.inform(core);
		try {
			synchronized (elevationCache) {
				elevationCache.clear();
				String f = initArgs.get(CONFIG_FILE);
				if (f == null) {
					throw new SolrException(
							SolrException.ErrorCode.SERVER_ERROR,
							"QueryElevationComponent must specify argument: '"
									+ CONFIG_FILE + "' -- path to elevate.xml");
				}
				boolean exists = false;

				// check if using ZooKeeper
				ZkController zkController = core.getCoreDescriptor()
						.getCoreContainer().getZkController();
				if (zkController != null) {
					// TODO : shouldn't have to keep reading the config name
					// when it has been read before
					exists = zkController.configFileExists(
							zkController.readConfigName(core
									.getCoreDescriptor().getCloudDescriptor()
									.getCollectionName()), f);
				} else {
					File fC = new File(core.getResourceLoader().getConfigDir(),
							f);
					File fD = new File(core.getDataDir(), f);
					if (fC.exists() == fD.exists()) {
						throw new SolrException(
								SolrException.ErrorCode.SERVER_ERROR,
								"QueryElevationComponent missing config file: '"
										+ f + "\n" + "either: "
										+ fC.getAbsolutePath() + " or "
										+ fD.getAbsolutePath()
										+ " must exist, but not both.");
					}
					if (fC.exists()) {
						exists = true;
						log.info("Loading QueryElevation from: "
								+ fC.getAbsolutePath());
						Config cfg = new Config(core.getResourceLoader(), f);
						elevationCache.put(null, loadElevationMap(cfg));
					}
				}
				// in other words, we think this is in the data dir, not the
				// conf dir
				if (!exists) {
					// preload the first data
					RefCounted<SolrIndexSearcher> searchHolder = null;
					try {
						searchHolder = core.getNewestSearcher(false);
						IndexReader reader = searchHolder.get()
								.getIndexReader();
						getElevationMap(reader, core);
					} finally {
						if (searchHolder != null)
							searchHolder.decref();
					}
				}
			}
		} catch (Exception ex) {
			throw new SolrException(SolrException.ErrorCode.SERVER_ERROR,
					"Error initializing QueryElevationComponent.", ex);
		}
	}

	// get the elevation map from the data dir
	protected Map<String, ElevationBean> getElevationMap(IndexReader reader,
			SolrCore core) throws Exception {
		synchronized (elevationCache) {
			Map<String, ElevationBean> map = elevationCache.get(null);
			if (map != null)
				return map;

			map = elevationCache.get(reader);
			if (map == null) {
				String f = initArgs.get(CONFIG_FILE);
				if (f == null) {
					throw new SolrException(
							SolrException.ErrorCode.SERVER_ERROR,
							"QueryElevationComponent must specify argument: "
									+ CONFIG_FILE);
				}
				log.info("Loading QueryElevation from data dir: " + f);

				Config cfg;

				ZkController zkController = core.getCoreDescriptor()
						.getCoreContainer().getZkController();
				if (zkController != null) {
					cfg = new Config(core.getResourceLoader(), f, null, null);
				} else {
					InputStream is = VersionedFile.getLatestFile(
							core.getDataDir(), f);
					cfg = new Config(core.getResourceLoader(), f,
							new InputSource(is), null);
				}

				map = loadElevationMap(cfg);
				elevationCache.put(reader, map);
			}
			return map;
		}
	}

	// load up the elevation map
	private Map<String, ElevationBean> loadElevationMap(Config cfg)
			throws IOException {
		XPath xpath = XPathFactory.newInstance().newXPath();
		Map<String, ElevationBean> map = new HashMap<String, ElevationBean>();
		NodeList nodes = (NodeList) cfg.evaluate("elevate/query",
				XPathConstants.NODESET);
		for (int i = 0; i < nodes.getLength(); i++) {
			Node node = nodes.item(i);
			String qstr = DOMUtil.getAttr(node, "text", "missing query 'text'");

			NodeList children = null;
			try {
				children = (NodeList) xpath.evaluate("doc", node,
						XPathConstants.NODESET);
			} catch (XPathExpressionException e) {
				throw new SolrException(SolrException.ErrorCode.SERVER_ERROR,
						"query requires '<doc .../>' child");
			}

			ArrayList<String> include = new ArrayList<String>();
			ArrayList<String> exclude = new ArrayList<String>();
			for (int j = 0; j < children.getLength(); j++) {
				Node child = children.item(j);
				String id = DOMUtil.getAttr(child, "id", "missing 'id'");
				String e = DOMUtil.getAttr(child, EXCLUDE, null);
				if (e != null) {
					if (Boolean.valueOf(e)) {
						exclude.add(id);
						continue;
					}
				}
				include.add(id);
			}

			ElevationBean elev = new ElevationBean(qstr, include, exclude);
			if (map.containsKey(elev.analyzed)) {
				throw new SolrException(SolrException.ErrorCode.SERVER_ERROR,
						"Boosting query defined twice for query: '" + elev.text
								+ "' (" + elev.analyzed + "')");
			}
			map.put(elev.analyzed, elev);
		}
		return map;
	}

}
